#ifndef ONNX_PARSER_H
#define ONNX_PARSER_H

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <spdlog/spdlog.h>

#include <fstream>
#include <memory>
#include <string>

#include "nvlogger.h"

namespace weilin {

class OnnxParse {
 public:
  OnnxParse(uint32_t flag =
                1U << static_cast<uint32_t>(
                    nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH),
            std::size_t workspaceSize = 1U << 20);

  void parse(const std::string& file);

  void save_model(const std::string& file);

  std::shared_ptr<nvinfer1::IHostMemory> get_model() { return mModel; }

  bool is_parse_success() { return this->isParseSuccess; }

 private:
  bool isParseSuccess;
  std::unique_ptr<nvinfer1::IBuilder> mBuilder;
  std::unique_ptr<nvinfer1::INetworkDefinition> mNetwork;
  std::unique_ptr<nvonnxparser::IParser> mParser;
  std::shared_ptr<nvinfer1::IBuilderConfig> mConfig;
  std::shared_ptr<nvinfer1::IHostMemory> mModel;
};

}  // namespace weilin

#endif  // !ONNX_PARSER_H
